from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
import numpy as np

np.random.seed(0)

def loadDatas():
    iris = load_iris()
    X = iris.data
    Y = iris.target
    X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.2)
    return X_train, X_test, y_train, y_test

X_train, X_test, y_train, y_test = loadDatas()


import matplotlib.pyplot as plt
def plotSamplesPoints(X, Y):
    """
    :param X: 形状为（N，2）的数组
    :param Y: 形状为（N，）的包含对应样本的类别标签的数组
    :return: None
    """
    # 颜色映射，为每个类别分配一个颜色
    colors = ['r', 'g', 'b']  # 假设有3个类别，分别用红、绿、蓝表示
    class_names = ['setosa', 'versicolor', 'virginica']

    # 绘制散点图
    plt.figure(figsize=(8, 6))  # 设置图形大小
    for i in range(3):
        # 筛选当前类别的数据点
        current_class_points = X[Y == i]
        plt.scatter(current_class_points[:, 0], current_class_points[:, 1], color=colors[i], label=class_names[i])

    plt.xlabel('Feature 1')  # 横轴标签
    plt.ylabel('Feature 2')  # 纵轴标签
    plt.title('Scatter Plot of Data Points by Class')  # 图形标题
    plt.legend()  # 显示图例
    plt.show()

# 可视化测试数据
plotSamplesPoints(X_test[:,[1,3]],y_test)

def experiment_KNN(k=5):
    model = KNeighborsClassifier(n_neighbors=k)
    model.fit(X_train, y_train)

    trainScore = model.score(X_train, y_train)
    testScore = model.score(X_test, y_test)
    print("KNN experiment with 'k={}' ".format(k))
    print(f'\t trainScore:{trainScore} \n\t testScore:{testScore}\n')

    y_test_predict = model.predict(X_test)
    plotSamplesPoints(X_test[:, [1, 3]], y_test_predict)

for k in [1, 3, 5, 7, 9, 20, 30]:
    experiment_KNN(k)
